'''
Created on 21. tra. 2011.

@author: Ana Banana
'''
from numpy import *;
import math, re;

def get_similarity_matrix (file):
    if isinstance(file, str):
        file = open(file, 'r');
    return loadtxt(file, int)

def get_score (sim_value, treshold, fun_type):
    if sim_value < treshold:
        return 0;
    if fun_type == 'log':
        return (double(math.log(sim_value, 10))/2)
       
    elif fun_type == 'logistic':
        return 1 / (1 + math.exp(-(sim_value-50)/10));
        
    else:
        return 1;

'''
Computes the cumulative score matrix using the defined
threshold. 
The fun_type can be log (still not supported) or score, which 
means that every value below threshold will be zero, and every
value above will account for one. 
Other matrices must have the same rows species as the matrix, and
will be used to compute the cumulative score matrix
Function also returns the number of accumulated columns
so one can calculate the cutoff value.
'''
def cumulative_score_matrix (treshold, fun_type, *matrices):
    hit_num = 0;
    col_num = 0;
    for mat in matrices:
        sz = size(mat,0);
        break;
    score_vec = zeros((sz,1));
    for i in range(len(score_vec)):
        score_vec[i] = 0;
    for mat in matrices:
        col_num += size(mat,1);
        for i in range (size(mat,0)):
            for j in range(size(mat,1)):
                score = get_score (mat[i,j], treshold, fun_type);
                #print "score:", score
                if (score != 0):
                    hit_num += 1;
                score_vec[i] += score;
                
    score_vec = score_vec / score_vec.max();
    return (score_vec, col_num, hit_num);

'''
Returns the estimated evolutionary core. 
perc - the maximum percentage of the score vector to
put into the core. The true estimated core can be 
much smaller, since the zero elements cannot be 
inside the core
'''
def estimated_evolutionary_core1 (score_vec, perc):
    if perc > 1:
        perc = double(perc / 100);
    # saves the original indices of the unsorted array
    # so one can index the original proteins
    indices = ndarray.argsort(score_vec, axis=0);
    #ndarray.sort(sc, axis=0);
    zero_els = 0;
    for i in range (len(sc)):
        if sc[i] == 0:
            zero_els += 1;
    print zero_els
    print len(score_vec)
    
    perc_z = double(zero_els) / len(score_vec);
    if 1-perc_z < perc:
        inds = range(zero_els, len(indices));
        to_return = indices[inds];
        to_ret_perc = 1-perc_z;
    else:
        lngth = int(len(sc)*perc);
        inds = range(len(indices)-lngth, len(indices));
        to_return = indices[inds];
        to_ret_perc = perc;
    return (to_return, to_ret_perc);
        
        
'''
Returns the estimated evolutionary core where all the
elements below the cutoff value are excluded
'''
def estimated_evolutionary_core2 (score_vec, cutoff):
    inds = [];
    for i in range(len(score_vec)): 
        if score_vec[i] >= cutoff:
            inds.append(i);
    return inds;

def generate_core_fasta (original_fasta, dest_fasta, indices):
    print indices
    if isinstance(original_fasta, str):
        original_fasta = open(original_fasta, 'r');
    if isinstance(dest_fasta, str):
        dest_fasta = open(dest_fasta, 'w');
    
    lines = original_fasta.readlines();
    for i in range(len(lines)):
        if i%2 == 0:
            br = re.findall('>P(.*?)(?:\s).*', lines[i])[0]
            br = [int(br)];
            print br
            if br in indices:
                dest_fasta.write(lines[i]);
                dest_fasta.write(lines[i+1]);
            
    dest_fasta.close();
    original_fasta.close();
    
    
mat = get_similarity_matrix(r'E:\workspace\Python\BioInfProjectSVN\temp\Celeg_Ecoli\result_matrix.txt');
mat1 = get_similarity_matrix(r'E:\workspace\Python\BioInfProjectSVN\temp\Celeg_Hpylo\result_matrix.txt');
mat2 = get_similarity_matrix(r'E:\workspace\Python\BioInfProjectSVN\temp\Celeg_Hsapi\result_matrix.txt');
(sc, col_num, hit_num) = cumulative_score_matrix(20, 'logistic', mat, mat1, mat2)

(res,p) = estimated_evolutionary_core1(sc, 1);
#print res
#print sc[res]
generate_core_fasta(r'E:\workspace\Python\BioInfProjectSVN\temp\Celeg_Ecoli\Celeg.faa', 'nja.faa', res);

